package com.java.compilation;

import com.java.compilation.helper.IDGenerator;
import com.java.compilation.helper.Keyword;
import com.java.compilation.helper.Kind;
import com.java.compilation.helper.TokenType;
import com.java.compilation.vm.Command;
import com.java.compilation.vm.Segment;

import java.io.IOException;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Set;

public class CompilationEngine_11 {
    private final JackTokenizer tokenizer;
    private final SymbolTable symbolTable;
    private final VMWriter vmWriter;
    private static final Set<String> opSymbols = new HashSet<>();
    private String className;

    public CompilationEngine_11(JackTokenizer tokenizer, SymbolTable symbolTable, VMWriter vmWriter) throws IOException {
        this.tokenizer = tokenizer;
        this.symbolTable = symbolTable;
        this.vmWriter = vmWriter;
        opSymbols.addAll(Arrays.asList("~", "+", "-", "*", "/", "&", "|", "<", ">", "="));

        readNextToken();
        // should class
        shouldTokenEquals(TokenType.KEYWORD, Keyword.CLASS.name());
        compileClass();
    }

    public void compileClass() throws IOException {
        // 'class' className '{' classVarDec* subroutineDec* '}'

        // className 标识符
        readNextToken();
        shouldTokenEquals(TokenType.IDENTIFIER);
        this.className = tokenizer.identifier();

        // '{'
        readNextToken();
        shouldTokenEquals(TokenType.SYMBOL, "{");

        // classVarDec* or subroutineDec*
        while (tokenizer.hasMoreTokens()) {
            readNextToken();
            if (tokenizer.tokenType() == TokenType.KEYWORD) {
                // classVarDec static field
                if (Keyword.STATIC == tokenizer.keyword() || Keyword.FIELD == tokenizer.keyword()) {
                    compileClassVarDec();
                } else if (Keyword.CONSTRUCTOR == tokenizer.keyword() || Keyword.FUNCTION == tokenizer.keyword()
                        || Keyword.METHOD == tokenizer.keyword()) {
                    // subroutineDec
                    compileSubroutine();
                }
            } else {
                // '}'
                shouldTokenEquals(TokenType.SYMBOL, "}");
            }
        }
    }

    public void compileClassVarDec() throws IOException {
        // ('static'|'field') type varName (',' varName)* ';'
        Kind kind = Kind.valueOf(tokenizer.keyword().name());

        // type
        readNextToken();
        String type = getType();

        // varName
        readNextToken();
        shouldTokenEquals(TokenType.IDENTIFIER);
        String name = tokenizer.identifier();
        symbolTable.define(name, type, kind);

        readNextToken();
        while (tokenizer.tokenType() == TokenType.SYMBOL && ",".equalsIgnoreCase(tokenizer.symbol())) {
            // ','
            // varName
            readNextToken();
            shouldTokenEquals(TokenType.IDENTIFIER);
            name = tokenizer.identifier();
            symbolTable.define(name, type, kind);

            readNextToken();
        }

        // ';'
        shouldTokenEquals(TokenType.SYMBOL, ";");
    }

    public void compileSubroutine() throws IOException {
        symbolTable.startSubroutine();
        // ('constructor'|'function'|'method'') ('void'|type) subroutineName '(' parameterList ')' subroutineBody
        shouldTokenEquals(TokenType.KEYWORD);
        Keyword keyword = tokenizer.keyword();

        // ('void'|type)
        readNextToken();
//        shouldTokenEquals(TokenType.KEYWORD);
//        Keyword type = tokenizer.keyword();

        // subroutineName
        readNextToken();
        shouldTokenEquals(TokenType.IDENTIFIER);
        String subroutineName = tokenizer.identifier();

        // '('
        readNextToken();
        shouldTokenEquals(TokenType.SYMBOL, "(");
        readNextToken();
        compileParameterList();
        // ')'
        shouldTokenEquals(TokenType.SYMBOL, ")");

        // subroutineBody
        // '{' varDec* statements '}'
        readNextToken();
        shouldTokenEquals(TokenType.SYMBOL, "{");

        // 读一个token，
        readNextToken();
        while (tokenizer.tokenType() == TokenType.KEYWORD && Keyword.VAR == tokenizer.keyword()) {
            compileVarDec();
            readNextToken();
        }

        createFunction(keyword, subroutineName);

        compileStatements();

        shouldTokenEquals(TokenType.SYMBOL, "}");
    }

    private void createFunction(Keyword keyword, String subroutineName) {
        int varCount = symbolTable.varCount(Kind.VAR);
        vmWriter.writeFunction(this.className + "." + subroutineName, varCount);

        if (keyword == Keyword.CONSTRUCTOR) {
            // push constant filedSize
            // call Memory.alloc 1
            // pop pointer 0
            vmWriter.writePush(Segment.CONST, symbolTable.varCount(Kind.FIELD));
            vmWriter.writeCall("Memory.alloc", 1);
            vmWriter.writePop(Segment.POINTER, 0);
        } else if (keyword == Keyword.METHOD) {
            // 设置this段
            vmWriter.writePush(Segment.ARG, 0);
            vmWriter.writePop(Segment.POINTER, 0);
        }
    }

    public void compileParameterList() throws IOException {

        // (((type varName)(',' type varName))*)?
        while (tokenizer.tokenType() != TokenType.SYMBOL || !")".equalsIgnoreCase(tokenizer.symbol())) {
            // type
            String type = tokenizer.keyword().name().toLowerCase();

            // varName
            readNextToken();
            shouldTokenEquals(TokenType.IDENTIFIER);
            String name = tokenizer.identifier();
            symbolTable.define(name, type, Kind.ARG);

            // 读下一个token
            readNextToken();
            if (tokenizer.tokenType() == TokenType.SYMBOL && ",".equalsIgnoreCase(tokenizer.symbol())) {
                // ','
                readNextToken();
            }
        }
    }

    public void compileVarDec() throws IOException {
        // 'var' type varName (',' varName)* ';'
        shouldTokenEquals(TokenType.KEYWORD);

        // type
        readNextToken();
//        shouldTokenEquals(TokenType.KEYWORD);
        String type = getType();

        // varName
        readNextToken();
        shouldTokenEquals(TokenType.IDENTIFIER);
        symbolTable.define(tokenizer.identifier(), type, Kind.VAR);

        readNextToken();
        while (tokenizer.tokenType() == TokenType.SYMBOL && ",".equalsIgnoreCase(tokenizer.symbol())) {
            // varName
            readNextToken();
            shouldTokenEquals(TokenType.IDENTIFIER);
            symbolTable.define(tokenizer.identifier(), type, Kind.VAR);

            readNextToken();
        }

        // ';'
        shouldTokenEquals(TokenType.SYMBOL, ";");
    }

    private String getType() {
        String type = "";
        if (tokenizer.tokenType() == TokenType.KEYWORD) {
            type = tokenizer.keyword().name().toLowerCase();
        } else if (tokenizer.tokenType() == TokenType.IDENTIFIER) {
            type = tokenizer.identifier();
        }
        return type;
    }

    public void compileStatements() throws IOException {
        // statement*
        while (true) {
            if (tokenizer.tokenType() == TokenType.KEYWORD && Keyword.LET == tokenizer.keyword()) {
                // letStatement
                compileLet();
            } else if (tokenizer.tokenType() == TokenType.KEYWORD && Keyword.IF == tokenizer.keyword()) {
                // ifStatement
                compileIf();
                continue;
            } else if (tokenizer.tokenType() == TokenType.KEYWORD && Keyword.WHILE == tokenizer.keyword()) {
                // whileStatement
                compileWhile();
            } else if (tokenizer.tokenType() == TokenType.KEYWORD && Keyword.DO == tokenizer.keyword()) {
                // doStatement
                compileDo();
            } else if (tokenizer.tokenType() == TokenType.KEYWORD && Keyword.RETURN == tokenizer.keyword()) {
                // returnStatement
                compileReturn();
            } else {
                break;
            }
            readNextToken();
        }
    }

    public void compileDo() throws IOException {
        // 'do' subroutineCall ';'
        // subroutineName
        readNextToken();
        shouldTokenEquals(TokenType.IDENTIFIER);
        String subroutineName = tokenizer.identifier();

        readNextToken();
        subroutineCall(subroutineName);

        // ';'
        readNextToken();
        shouldTokenEquals(TokenType.SYMBOL, ";");
    }

    public void compileLet() throws IOException {
        // 'let' varName ('['expression']')?'='expression';'
        shouldTokenEquals(TokenType.KEYWORD);

        // varName
        readNextToken();
        shouldTokenEquals(TokenType.IDENTIFIER);
        String varName = tokenizer.identifier();

        readNextToken();
        // 数组处理
        if (tokenizer.tokenType() == TokenType.SYMBOL && "[".equalsIgnoreCase(tokenizer.symbol())) {
            readNextToken();
            compileExpression();
            // ']'
//            readNextToken();
            shouldTokenEquals(TokenType.SYMBOL, "]");

            readNextToken();
        }

        // '='
        shouldTokenEquals(TokenType.SYMBOL);

        readNextToken();
        // expression
        compileExpression();

//        readNextToken();
        shouldTokenEquals(TokenType.SYMBOL, ";");

        // pop varName
        pop(varName);
    }

    private void push(String varName) {
        Kind kind = symbolTable.kindOf(varName);
        if (kind == Kind.ARG) {
            vmWriter.writePush(Segment.ARG, symbolTable.indexOf(varName));
        } else if (kind == Kind.FIELD) {
            vmWriter.writePush(Segment.THIS, symbolTable.indexOf(varName));
        } else if (kind == Kind.VAR) {
            vmWriter.writePush(Segment.LOCAL, symbolTable.indexOf(varName));
        } else if (kind == Kind.STATIC) {
            vmWriter.writePush(Segment.STATIC, symbolTable.indexOf(varName));
        }
    }

    private void pop(String varName) {
        Kind kind = symbolTable.kindOf(varName);
        if (kind == Kind.ARG) {
            vmWriter.writePop(Segment.ARG, symbolTable.indexOf(varName));
        } else if (kind == Kind.FIELD) {
            vmWriter.writePop(Segment.THIS, symbolTable.indexOf(varName));
        } else if (kind == Kind.VAR) {
            vmWriter.writePop(Segment.LOCAL, symbolTable.indexOf(varName));
        } else if (kind == Kind.STATIC) {
            vmWriter.writePop(Segment.STATIC, symbolTable.indexOf(varName));
        }
    }

    public void compileWhile() throws IOException {
        int whileLabel1 = IDGenerator.getAndIncrLabel();
        int whileLabel2 = IDGenerator.getAndIncrLabel();
        // 'while' '(' expression ')' '{' statements '}'
        shouldTokenEquals(TokenType.KEYWORD, Keyword.WHILE.name());

        vmWriter.writeLabel("while_" + whileLabel1);
        // '('
        readNextToken();
        shouldTokenEquals(TokenType.SYMBOL, "(");
        readNextToken();
        compileExpression();
        // ')'
        shouldTokenEquals(TokenType.SYMBOL, ")");

        vmWriter.writeIf("while_" + whileLabel2);
        vmWriter.writeGoto("while_" + whileLabel1 + "_end");
        // '{'
        vmWriter.writeLabel("while_" + whileLabel2);
        readNextToken();
        shouldTokenEquals(TokenType.SYMBOL, "{");

        readNextToken();
        compileStatements();
        vmWriter.writeGoto("while_" + whileLabel1);

        // '}'
        shouldTokenEquals(TokenType.SYMBOL, "}");
        vmWriter.writeLabel("while_" + whileLabel1 + "_end");
    }

    public void compileReturn() throws IOException {
        // 'return' expression ';'

        readNextToken();
        if (tokenizer.tokenType() != TokenType.SYMBOL || !";".equalsIgnoreCase(tokenizer.symbol())) {
            compileExpression();
        } else {
            vmWriter.writePush(Segment.CONST, 0);
        }

        vmWriter.writeReturn();
        // ';'
        shouldTokenEquals(TokenType.SYMBOL, ";");
    }

    public void compileIf() throws IOException {
        int ifLabel = IDGenerator.getAndIncrLabel();
        // 'if' '('expression')''{'statements'}'
        // ('else' '{'statements'}')?
        shouldTokenEquals(TokenType.KEYWORD, Keyword.IF.name());
        // '('
        readNextToken();
        readNextToken();
        compileExpression();
        // ')'
        vmWriter.writeIf("if_" + ifLabel);
        vmWriter.writeGoto("if_" + ifLabel + "_else");

        // '{'
        readNextToken();
        shouldTokenEquals(TokenType.SYMBOL, "{");

        vmWriter.writeLabel("if_" + ifLabel);
        readNextToken();
        compileStatements();
        vmWriter.writeGoto("if_" + ifLabel + "_end");
        // '}'
        shouldTokenEquals(TokenType.SYMBOL, "}");

        readNextToken();
        // 'else'
        if (tokenizer.tokenType() == TokenType.KEYWORD && Keyword.ELSE == tokenizer.keyword()) {
            shouldTokenEquals(TokenType.KEYWORD, Keyword.ELSE.name());
            vmWriter.writeLabel("if_" + ifLabel + "_else");
            // '{'
            readNextToken();
            shouldTokenEquals(TokenType.SYMBOL, "{");

            readNextToken();
            compileStatements();

            shouldTokenEquals(TokenType.SYMBOL, "}");

            readNextToken();
        } else {
            vmWriter.writeLabel("if_" + ifLabel + "_else");
        }

        vmWriter.writeLabel("if_" + ifLabel + "_end");
    }

    public void compileExpression() throws IOException {
        // term (op term)*
        compileTerm();

        // (op term)*
        while (tokenizer.tokenType() == TokenType.SYMBOL && opSymbols.contains(tokenizer.symbol())) {
            // op
            shouldTokenEquals(TokenType.SYMBOL);
            String symbol = tokenizer.symbol();

            readNextToken();
            compileTerm();

            if ("*".equalsIgnoreCase(symbol)) {
                vmWriter.writeCall("Math.multiply", 2);
            } else if ("/".equalsIgnoreCase(symbol)) {
                vmWriter.writeCall("Math.divide", 2);
            } else {
                vmWriter.writeArithmetic(Command.valueFrom(symbol));
            }
        }
    }

    public void compileTerm() throws IOException {
        // integerConstant | stringConstant | keywordConstant | varName | varName '[' expression ']' | subroutineCall |
        // '('expression ')' | unaryOp term
        //  integerConstant | stringConstant | keywordConstant
        if (tokenizer.tokenType() == TokenType.INT_CONST) {
            vmWriter.writePush(Segment.CONST, tokenizer.intVal());
            readNextToken();
            return;
        } else if (tokenizer.tokenType() == TokenType.STRING_CONST) {
            System.out.println("not done");
            readNextToken();
            return;
        } else if (tokenizer.tokenType() == TokenType.KEYWORD) {
            if (Keyword.TRUE == tokenizer.keyword()) {
                vmWriter.writePush(Segment.CONST, 0);
                vmWriter.writeArithmetic(Command.NOT);
            } else if (Keyword.FALSE == tokenizer.keyword()) {
                vmWriter.writePush(Segment.CONST, 0);
            } else if (Keyword.THIS == tokenizer.keyword()) {
                vmWriter.writePush(Segment.POINTER, 0);
            }
            readNextToken();
            return;
        }

        // '('expression ')'
        if (tokenizer.tokenType() == TokenType.SYMBOL && "(".equalsIgnoreCase(tokenizer.symbol())) {
            // '('
            shouldTokenEquals(TokenType.SYMBOL, "(");
            readNextToken();
            compileExpression();
            // ')'
            shouldTokenEquals(TokenType.SYMBOL, ")");
            readNextToken();
            return;
        }

        // varName | varName '[' expression ']' | subroutineCall
        if (tokenizer.tokenType() == TokenType.IDENTIFIER) {
            shouldTokenEquals(TokenType.IDENTIFIER);
            String varName = tokenizer.identifier();

            readNextToken();
            if (tokenizer.tokenType() == TokenType.SYMBOL && "[".equalsIgnoreCase(tokenizer.symbol())) {
                // '['
                readNextToken();
                // 数组处理
                compileExpression();

                // ']'
                shouldTokenEquals(TokenType.SYMBOL, "]");

                readNextToken();
            } else if (tokenizer.tokenType() == TokenType.SYMBOL
                    && ("(".equalsIgnoreCase(tokenizer.symbol())
                    || ".".equalsIgnoreCase(tokenizer.symbol()))) {
                subroutineCall(varName);

                readNextToken();
            } else {
                // varName
                push(varName);
            }

            return;
        }

        if (tokenizer.tokenType() == TokenType.SYMBOL) {
            String symbol = tokenizer.symbol();

            readNextToken();
            compileTerm();
            if ("-".equalsIgnoreCase(symbol)) {
                vmWriter.writeArithmetic(Command.NEG);
            } else if ("~".equalsIgnoreCase(symbol)) {
                vmWriter.writeArithmetic(Command.NOT);
            }
        }
    }

    private void subroutineCall(String subroutineName) throws IOException {
        int argCount = 0;
        // subroutineName '(' expressionList ')' | (className | varName)'.' subroutineName '(' expressionList ')'
        StringBuilder callee = new StringBuilder(subroutineName);
        if (tokenizer.tokenType() == TokenType.SYMBOL && ".".equalsIgnoreCase(tokenizer.symbol())) {
            // (className | varName)'.' subroutineName '(' expressionList ')'
            if (symbolTable.typeOf(subroutineName) != null) {
                callee = new StringBuilder(symbolTable.typeOf(subroutineName));
                // method首个参数是this
                push(subroutineName);
                argCount++;
            }

            // '.'
            callee.append(tokenizer.symbol());
            // subroutineName
            readNextToken();
            shouldTokenEquals(TokenType.IDENTIFIER);
            callee.append(tokenizer.identifier());

            readNextToken();
        } else {
            callee = new StringBuilder(className + "." + callee);
            // subroutineName '(' expressionList ')'
            // method首个参数是this
            vmWriter.writePush(Segment.POINTER, 0);
            argCount++;
        }
        // '('
        shouldTokenEquals(TokenType.SYMBOL, "(");
        // expressionList
        argCount += compileExpressionList();
        // ')'
        shouldTokenEquals(TokenType.SYMBOL, ")");
        vmWriter.writeCall(callee.toString(), argCount);
    }

    public int compileExpressionList() throws IOException {
        int argCount = 0;
        readNextToken();
        if (tokenizer.tokenType() == TokenType.SYMBOL && ")".equalsIgnoreCase(tokenizer.symbol())) {
            return argCount;
        }

        argCount++;
        // (expression(','expression)*)?
        compileExpression();
        // (','expression)*
        while (tokenizer.tokenType() == TokenType.SYMBOL && ",".equalsIgnoreCase(tokenizer.symbol())) {
            readNextToken();
            compileExpression();
            argCount++;
        }
        return argCount;
    }

    private void shouldTokenEquals(TokenType tokenType) {
        if (tokenizer.tokenType() != tokenType) {
            throw new IllegalArgumentException(String.format("token type not equals %s", tokenType));
        }
    }

    private void shouldTokenEquals(TokenType tokenType, String value) {
        shouldTokenEquals(tokenType);

        if (tokenType == TokenType.KEYWORD && !tokenizer.keyword().name().equalsIgnoreCase(value)) {
            throw new IllegalArgumentException(String.format("keyword not equals %s", value));
        } else if (tokenType == TokenType.SYMBOL && !value.equalsIgnoreCase(tokenizer.symbol())) {
            throw new IllegalArgumentException(String.format("symbol not equals %s", value));
        }
    }

    private void readNextToken() throws IOException {
        if (!tokenizer.hasMoreTokens()) {
            throw new IllegalArgumentException("read token error");
        }
        tokenizer.advance();
    }


}
